package com.imooc;

import com.alibaba.csp.sentinel.Entry;
import com.alibaba.csp.sentinel.SphU;
import com.alibaba.csp.sentinel.context.ContextUtil;
import com.alibaba.csp.sentinel.slots.block.BlockException;
import com.alibaba.csp.sentinel.slots.block.RuleConstant;
import com.alibaba.csp.sentinel.slots.block.flow.FlowRule;
import com.alibaba.csp.sentinel.slots.block.flow.FlowRuleManager;
import com.alibaba.csp.sentinel.util.TimeUtil;

import java.util.ArrayList;
import java.util.List;
import java.util.Random;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @Description
 * @Author 顾志杰
 * @Date 2021/9/26 16:04
 */
public class FlowQpsDemo {
    // 资源名称
    private static final String KEY = "abc";
    // 统计执行通过的总数
    private static AtomicInteger pass = new AtomicInteger();
    // 统计执行限制的总数
    private static AtomicInteger block = new AtomicInteger();
    // 统计所有执行总数
    private static AtomicInteger total = new AtomicInteger();

    private static volatile boolean stop = false;
    // 启动的线程数
    private static final int threadCount = 32;
    // 执行时间单位为秒(s)
    private static int seconds = 60 + 40;

    public static void main(String[] args) throws Exception {
        // 初始化流控规则
        initFlowQpsRule();
        // 启动统计线程
        tick();
        // 启动业务线程
        simulateTraffic();

        System.out.println("===== begin to do flow control");
        System.out.println("only 20 requests per second can pass");

    }

    private static void initFlowQpsRule() {
        List<FlowRule> rules = new ArrayList<FlowRule>();
        FlowRule rule1 = new FlowRule();
        // 资源名，即限流规则的作用对象
        rule1.setResource(KEY);
        // set limit qps to 20 限流阈值
        rule1.setCount(20);
        // 限流阈值类型（QPS 或并发线程数）
        rule1.setGrade(RuleConstant.FLOW_GRADE_QPS);
        // 流控针对的调用来源，若为 default 则不区分调用来源
        rule1.setLimitApp("default");
        // 调用关系限流策略
        // rule1.setStrategy(RuleConstant.STRATEGY_DIRECT);
        // 流量控制效果（直接拒绝、Warm Up、匀速排队）
        // rule1.setControlBehavior(RuleConstant.CONTROL_BEHAVIOR_DEFAULT);
        rules.add(rule1);
        // 加载限流的规则
        FlowRuleManager.loadRules(rules);
    }

    private static void simulateTraffic() {
        for (int i = 0; i < threadCount; i++) {
            Thread t = new Thread(new RunTask());
            t.setName("simulate-traffic-Task");
            t.start();
        }
    }

    private static void tick() {
        Thread timer = new Thread(new TimerTask());
        timer.setName("sentinel-timer-task");
        timer.start();
    }

    static class TimerTask implements Runnable {

        @Override
        public void run() {
            long start = System.currentTimeMillis();
            System.out.println("begin to statistic!!!");

            long oldTotal = 0;
            long oldPass = 0;
            long oldBlock = 0;
            while (!stop) {
                try {
                    TimeUnit.SECONDS.sleep(1);
                } catch (InterruptedException e) {
                }
                long globalTotal = total.get();
                long oneSecondTotal = globalTotal - oldTotal;
                oldTotal = globalTotal;

                long globalPass = pass.get();
                long oneSecondPass = globalPass - oldPass;
                oldPass = globalPass;

                long globalBlock = block.get();
                long oneSecondBlock = globalBlock - oldBlock;
                oldBlock = globalBlock;

                System.out.println(seconds + " send qps is: " + oneSecondTotal);
                System.out.println(TimeUtil.currentTimeMillis() + ", total:" + oneSecondTotal
                        + ", pass:" + oneSecondPass
                        + ", block:" + oneSecondBlock);

                if (seconds-- <= 0) {
                    stop = true;
                }
            }

            long cost = System.currentTimeMillis() - start;
            System.out.println("time cost: " + cost + " ms");
            System.out.println("total:" + total.get() + ", pass:" + pass.get()
                    + ", block:" + block.get());
            System.exit(0);
        }
    }

    static class RunTask implements Runnable {
        @Override
        public void run() {
            while (!stop) {
                ContextUtil.enter("entrancel", "appA");
                try(Entry entry = SphU.entry(KEY)) {
                    // token acquired, means pass
                    pass.addAndGet(1);
                } catch (BlockException e1) {
                    block.incrementAndGet();
                } catch (Exception e2) {
                    // biz exception
                } finally {
                    total.incrementAndGet();
                }
                ContextUtil.exit();
                Random random2 = new Random();
                try {
                    TimeUnit.MILLISECONDS.sleep(random2.nextInt(50));
                } catch (InterruptedException e) {
                    // ignore
                }
            }
        }
    }
}
